/*
 * Project CHEN Rendering Engine : /sppm.cpp
 * Created: 2021-01-27, Last modified:  2021-02-04
 * 
 * This file implements the methods of the SPPMShader class and related
 * structs and functions.
 * 
 * Copyright (C) 2021 Hineven (hineven@pku.edu.cn)
 * 
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 * 
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 * 
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include "sppm.hpp"
namespace chen {

static int resolution_w;
static int resolution_h;

static VisiblePoint * vpoints;

static struct SPPMPixel {
    Float N, M;
    Spectrum G;
    Float radius;
    Spectrum Phi;
    int atrN; // attached number
    Spectrum atrV; // attached spectrums
} * image;

static struct TrSPPMPixel {
    int M;
    Spectrum G;
} ** pixels;
// Each thread owns its own pixel map to avoid race conditions.

struct VisiblePoint {
    Point3 p; // world position
    Spectrum f; // attenuation
    Vector3 wo;
    BSDF * bsdf; // interaction bsdf (cloned), including local tangent space info
    int pixel_id;
    Float pixel_rr; // pixel radius ^ 2 (cached here for acceleration)
};

struct KDTreeNode {
    int l, r;
    Bound3 bnd;
    VisiblePoint vp;
    Float mrr;
} ;

static int _cdim;
static inline bool dimcomp (const VisiblePoint & a, const VisiblePoint & b) {
    return a.p[_cdim] < b.p[_cdim];
}


class Packeter {
public: 
    pthread_mutex_t mut;
    int remain, pack;
    Packeter (int total, int per_pack): remain(total), pack(per_pack) {
        pthread_mutex_init(&mut, nullptr);
    }
    ~Packeter () {
        pthread_mutex_destroy(&mut);
    }
    int fetch () {
        printf("\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\bre:%d        ", remain);
        if(remain == 0) return 0;
        pthread_mutex_lock(&mut);
        int re = min(pack, remain);
        remain -= re;
        pthread_mutex_unlock(&mut);
        return re;
    }
} ;

static int iter_num;

struct KDTreeQuery { 
    // Save parameters for a K-D Tree query
    // These parameters are not required to be pushed into stack
    // as the recurision proceeds, thus accelerating the whole
    // process.
    Point3 p;
    Spectrum sc_photon;
    Vector3 wi;
    TrSPPMPixel * tr_pixels;
} ;

static struct KDTree {
    KDTreeNode * d;
    int dtot, root;
    VisiblePoint * target;
    KDTree (int npoints) {
        LOG(INFO) << "[SPPM] " << (npoints+1)*sizeof(KDTreeNode)/1024 << " KB memory allocated for K-D Tree."; 
        d = new KDTreeNode[npoints+1];
    }
    ~KDTree () {
        delete [] d;
    }
    inline void update (int u) {
        d[u].bnd = Bound3(d[u].vp.p);
        d[u].mrr = d[u].vp.pixel_rr;
        if(d[u].l) {
            d[u].bnd = combine(d[u].bnd, d[d[u].l].bnd);
            d[u].mrr = max(d[u].mrr, d[d[u].l].mrr);
        }
        if(d[u].r) {
            d[u].bnd = combine(d[u].bnd, d[d[u].r].bnd);
            d[u].mrr = max(d[u].mrr, d[d[u].r].mrr);
        }
    }
    int _build (int l, int r) {
        if(l > r) return 0;
        int u = ++dtot;
        d[u].l = d[u].r = 0;
        if(l == r) {
            d[u].vp = target[l];
            d[u].bnd = Bound3(target[l].p);
            d[u].mrr = target[l].pixel_rr;
            return u;
        }
        int mxdim = 0, mxv = 0;
        for(int i = 0; i<3; i++) {
            Float mn = Infinity;
            Float mx = -Infinity;
            for(int j = l; j<=r; j++) {
                mn = min(mn, target[j].p[i]);
                mx = max(mx, target[j].p[i]);
            }
            if(mxv < mx-mn) mxv = mx-mn, mxdim = i;
        }
        int mid = (l+r)>>1;
        _cdim = mxdim;
        std::nth_element(target+l, target+mid, target+r+1, dimcomp);
        d[u].vp = target[mid];
        d[u].l = _build(l, mid-1);
        d[u].r = _build(mid+1, r);
        update(u);
        return u;
    }
    void clear () {
        dtot = 0;
    }
    void build (VisiblePoint * points, int npoints) {
        target = points;
        root = _build(0, npoints-1);
    }
    // Automatically accumulate radiance to (the thread's) pixel map
    // This function should support parrallism
    int _query (int u, const KDTreeQuery * q) {
        if(!u) return 0;
        if(d[u].bnd.minEulerDistSquared(q->p) > d[u].mrr) return 0;
        int t = 0;
        if((d[u].vp.p-q->p).lengthSquared() <= d[u].vp.pixel_rr) {
            t = 1;
            // Accumulate photon to thread pixel map
            q->tr_pixels[d[u].vp.pixel_id].M ++;
            q->tr_pixels[d[u].vp.pixel_id].G += d[u].vp.bsdf->f(q->wi, d[u].vp.wo)*q->sc_photon*d[u].vp.f;
        }
        return _query(d[u].l, q) + t +_query(d[u].r, q);
    }
    int query (KDTreeQuery * q) {
        return _query(root, q);
    }
} * kdtree;

void SPPMShader::prepare (int ntr) {
    Shader::prepare(ntr);
    LOG(INFO) << "[SPPM] Preparing.";
    resolution_w = camera->film->resolution_w;
    resolution_h = camera->film->resolution_h;
    image = new SPPMPixel[resolution_h*resolution_w];
    LOG(INFO) << "[SPPM] " << resolution_h*resolution_w*sizeof(SPPMPixel)/1024 << " KB memory allocated for final image.";
    SPPMPixel zip;
    zip.G = 0;
    zip.N = 0;
    zip.radius = initial_radius;
    std::fill(image, image+resolution_w*resolution_h, zip);
    pixels = new TrSPPMPixel*[nthreads];
    int mt = 0;
    for(int i = 0; i<nthreads; i++) {
        pixels[i] = new TrSPPMPixel[resolution_h*resolution_w];
        if(pixels[i] == nullptr) LOG(FATAL) << "Memory fucked up creating copies of the pixel map.";
        TrSPPMPixel zp;
        zp.G = 0;
        zp.M = 0;
        std::fill(pixels[i], pixels[i]+resolution_h*resolution_w, zp);
        mt += resolution_h*resolution_w*sizeof(TrSPPMPixel);
    }
    LOG(INFO) << "[SPPM] A total of " << mt/1024 << " KB memory allocated for thread buffers.";
    vpoints = new VisiblePoint[resolution_h*resolution_w*samples_per_pixel];
    if(vpoints == nullptr) {
        LOG(FATAL) << "Memory drained creating visible points.";
    }
    for(int i = 0; i<resolution_h*resolution_w*samples_per_pixel; i++)
        vpoints[i].bsdf = nullptr;
    LOG(INFO) << "[SPPM] " << sizeof(VisiblePoint)*(resolution_h*resolution_w*samples_per_pixel)/1024
        << " KB memory allocated for visible points.";
    kdtree = new KDTree(resolution_h*resolution_w*samples_per_pixel);
}

SPPMShader::~SPPMShader () {
    delete [] image;
    for(int i = 0; i<nthreads; i++) {
        delete [] pixels[i];
    }
    delete [] pixels;
    delete [] vpoints;
    delete kdtree;
}

void SPPMShader::traceCameraRays () {
    resetFetchNext();
    std::thread**tr = new std::thread*[nthreads];
    for(int i = 0; i<nthreads; i++) {
tr[i] = new std::thread([this](int thread_id) {
    Sampler * sampler = sampler_f->clone(thread_id+iter_num*nthreads);
// Paralleled tracing camera rays
    int base_w, base_h;
    while(fetchNext(&base_w, &base_h)) {
        for(int dw = 0; base_w+dw<resolution_w && dw < BLOCK_SIZE; dw++)
            for(int dh = 0; base_h+dh<resolution_h && dh < BLOCK_SIZE; dh++)
                for(int i = 0; i<samples_per_pixel; i++) {
                    VisiblePoint * vp = &vpoints[(dw+base_w+(dh+base_h)*resolution_w)*samples_per_pixel+i];
                    if(vp->bsdf != nullptr) bfree(vp->bsdf), vp->bsdf = nullptr;
                    vp->pixel_id = dw+base_w+(dh+base_h)*resolution_w;
                    vp->pixel_rr = image[vp->pixel_id].radius*image[vp->pixel_id].radius;
                    Point2 u = sampler->gen2D();
                    Float pdf;
                    RayDifferential rd = camera->rayDifferentialFromCamera(dw+base_w+u.x, dh+base_h+u.y, &pdf);
                    if(!traceCameraRay(rd, pdf, sampler, vp)) {
                        vp->pixel_id = -1;
                    }
                }
    }
    delete sampler;
}, i);
    }
    for(int i = 0; i<nthreads; i++) {
        tr[i]->join();
        delete tr[i];
    }
    delete [] tr;
    int curp = 0;
    for(int i = 0; i<resolution_w*resolution_h*samples_per_pixel; i++)
        if(vpoints[i].pixel_id != -1) {
            vpoints[curp++] = vpoints[i]; 
            if(curp != i+1) vpoints[i].bsdf = nullptr;
        }
    // Construct datastructure
    kdtree->clear();
    kdtree->build(vpoints, curp);
}

static RayDifferential specularBounce (
    const RayDifferential & rd, const SurfInteract & si,
    const Vector3 & nd, unsigned sampled_tags, Float * eta
) {
    // used to calculate the direction
    bool reflect = hasFlag(sampled_tags, FBSDF_REFLECTION);
    RayDifferential ret = si.spawnRayPickSide(nd);
    if(reflect) {
        // Specular reflection
        Normal n = si.ns;
        Vector3 nx = ((Vector3)n+(Vector3)si.dndx).normalized();
        Vector3 ny = ((Vector3)n+(Vector3)si.dndy).normalized();
        ret.dxori = si.p+si.dpdx;
        ret.dyori = si.p+si.dpdy;
        ret.dxdir = rd.dxdir-(Float)2*dot(rd.dxdir, nx)*nx;
        ret.dydir = rd.dxdir-(Float)2*dot(rd.dydir, ny)*ny;
    } else {
        // Specular transmission
        Normal n = si.ns;
        Normal nx = (Normal)((Vector3)n+(Vector3)si.dndx).normalized();
        Normal ny = (Normal)((Vector3)n+(Vector3)si.dndy).normalized();
        ret.dxori = si.p+si.dpdx;
        ret.dyori = si.p+si.dpdy;
        Normal nxy = (Normal)nx.getPerpendicularUnit();
        // Fresnel specular transmission direction is isotropic respect 
        // to coordinate system. So temporarily generated system is valid.
        ret.dxdir = nx.rcoord(
            nxy,
            Fresnel(si.bsdf->eta).transmitDirection(
                nx.coord(nxy, -rd.dxdir)
            )
        );
        Normal nyy = (Normal)ny.getPerpendicularUnit();
        ret.dydir = ny.rcoord(
            nyy,
            Fresnel(si.bsdf->eta).transmitDirection(
                ny.coord(nyy, -rd.dydir)
            )
        );
    }
    // update eta
    if(!si.ng.sideWith(rd.d) && si.ng.sideWith(nd)) { // entering
        *eta *= si.bsdf->eta;
    } else if(si.ng.sideWith(rd.d) && !si.ng.sideWith(nd)) {
        *eta /= si.bsdf->eta;
    }
    return ret;
}

bool SPPMShader::traceCameraRay (RayDifferential r, Float crpdf, Sampler * sampler, VisiblePoint * vp) {
    if(crpdf == 0) return false;
    SurfInteract si;
    Float pdf;
    Vector3 wi;
    Spectrum beta = 1;
    int bounces = 0;
    Spectrum ret = 0;
    Float cur_eta = 1.0;
    bool specular_bounce = false;
    while(true) {
        if(!scene->castRay(r, &si)) {
            Spectrum sum = 0;
            for(int i = 0; i<(int)scene->lights.size(); i++)
                sum += scene->lights[i]->LeFrom(r);
            ret += sum*beta;
            image[vp->pixel_id].atrN ++;
            image[vp->pixel_id].atrV += ret/crpdf;
            return false;
        }
        si.calculateDifferntials(r);
        si.applyMaterial(false);
        if(bounces == 0 || specular_bounce) {
            ret += beta*si.Le();
            specular_bounce = false;
        }
        if(si.bsdf != nullptr) {
            if(bounces <= max_bounce && si.bsdf->count(FBSDF_DELTA)) {
                unsigned tags;
                Spectrum f = si.bsdf->sample(-si.wh, sampler->gen2D(), &wi, &pdf, true, &tags);
                if(pdf == 0 || f == 0) return false;
                beta *= f*(abs(si.ns.weigh(wi))/pdf);
                // Specular surface, need to keep ray differential
                if(hasFlag(tags, FBSDF_SPECULAR)) {
                    specular_bounce = true;
                    r = specularBounce(r, si, -wi, tags, &cur_eta);
                }
            } else {
                ret += beta*estimateDirectIlluminationND(si, nlight_samples, sampler);
                break ;
            }
        } else {
            // null intersection encountered, continue going
            r = si.continueRayDiff(r, Infinity);
        }
        if(bounces > max_bounce) return false;
        bounces ++;
    }
    image[vp->pixel_id].atrN ++;
    image[vp->pixel_id].atrV += ret/crpdf;
    vp->p = si.p;
    // Clone the bsdf function
    vp->bsdf = si.bsdf->clone();
    vp->f = beta/crpdf;
    vp->wo = -si.wh;
    return true;
}

void SPPMShader::tracePhotons () {
    // Paralleled tracing photons
    Packeter * pak = new Packeter(photons_per_iter, 1024);
    std::thread**tr = new std::thread*[nthreads];
    for(int i = 0; i<nthreads; i++) {
tr[i] = new std::thread([this, pak](int thread_id) {
    Sampler * sampler = sampler_f->clone(thread_id+iter_num*nthreads);
    int batch;
    while((batch = pak->fetch())) {
        while(batch --) traceRandomPhoton(thread_id, sampler);
    }
    delete sampler;
}, i);
    }
    for(int i = 0; i<nthreads; i++) {
        tr[i]->join();
        delete tr[i];
    }
    delete [] tr;
    delete pak;
}

void SPPMShader::traceRandomPhoton (int tid, Sampler * sampler) {
    Float lpdf;
    const Light * lemit = scene->pickLightRespectToArea(sampler->gen(), &lpdf);
    if(lemit == nullptr || !hasFlag(lemit->flags, FLIGHT_SAMPLE_LO)) return ;
    Ray r;
    Float pdf, cosv;
    // note that this time cosv is used to calculate photon powers
    Spectrum beta = lemit->sampleLo(sampler->gen2D(), sampler->gen2D(), &r, &pdf, &cosv);
    beta = beta*cosv/(pdf*lpdf);
    if(beta == 0) return ;
    SurfInteract si;
    int bounces = 0;
    while(true) {
        if(!scene->castRay(r, &si)) return ;
        if(bounces > 0) { // account for indirect illumination only
            KDTreeQuery kdq;
            kdq.p = si.p;
            kdq.wi = si.wh;
            kdq.sc_photon = beta;
            kdq.tr_pixels = pixels[tid];
            kdtree->query(&kdq);
        }
        if(bounces > max_bounce) return ;
        bounces ++;
        si.applyMaterial(false);
        Vector3 wi;
        Float dpdf;
        Spectrum f = si.bsdf->sample(si.wh, sampler->gen2D(), &wi, &dpdf, false, nullptr);
        if(f == 0) return ;
        r = si.spawnRayPickSide(-wi);
        Spectrum nbeta = beta*f*abs(si.ns.weigh(wi))/dpdf;
        if(bounces > 3) {
            Float q = max((Float).05, (Float)1-beta.Y()/nbeta.Y());
            if(rng.rand01() < q) return ; // thread local random number generator
            beta = nbeta/((Float)1-q);
        }
    }
}


void SPPMShader::gatherResults () {
    // Single thread gathering for simplicity
    for(int i = 0; i<resolution_h; i++)
        for(int j = 0; j<resolution_w; j++) {
            image[i*resolution_w+j].M = 0;
            image[i*resolution_w+j].Phi = 0;
        }
    for(int tr = 0; tr<nthreads; tr ++)
        for(int i = 0; i<resolution_h; i++)
            for(int j = 0; j<resolution_w; j++) {
                image[i*resolution_w+j].M += pixels[tr][i*resolution_w+j].M;
                image[i*resolution_w+j].Phi += pixels[tr][i*resolution_w+j].G;
            }
    TrSPPMPixel npixel;
    npixel.M = 0;
    npixel.G = 0;
    for(int i = 0; i<nthreads; i++)
        std::fill(pixels[i], pixels[i]+resolution_w*resolution_h, npixel);
    camera->film->clear();
    Sample * output = new Sample[resolution_w*resolution_h];
    Float avgR = 0, maxR = 0, minR = Infinity;
    Float Nall = (Float)photons_per_iter*iter_num;
    for(int i = 0; i<resolution_h; i++)
        for(int j = 0; j<resolution_w; j++) {
            SPPMPixel & p = image[i*resolution_w+j];
            Float aM = p.M*alpha;
            if(p.N + p.M) p.radius = p.radius*std::sqrt((Float)(p.N+aM)/(p.N+p.M));
            if(p.N + p.M) p.G = (p.G+p.Phi/samples_per_pixel)*((Float)(p.N+aM)/(p.N+p.M));
            p.N += aM;
            output[i*resolution_w+j].p = Point2(j+(Float)0.5, i+(Float)0.5);
            if(Nall) output[i*resolution_w+j].s = /*p.G/(Nall*p.radius*p.radius*(Float)M_PI)+*/zdiv(p.atrV, p.atrN);
            else output[i*resolution_w+j].s = zdiv(p.atrV, p.atrN);
            avgR += p.radius;
            maxR = max(maxR, p.radius);
            minR = min(minR, p.radius);
        }
    camera->submit(output, resolution_w*resolution_h);
    LOG(INFO) << "[SPPM] Average radius: " << avgR/(resolution_h*resolution_w) << " (" << minR << ", " << maxR << ")"; 
    delete [] output;
}

void SPPMShader::render () {
    for(iter_num = 1; iter_num <= iters; iter_num ++) {
        LOG(INFO) << "[SPPM] Iteration #" << iter_num;
        traceCameraRays();
        LOG(INFO) << "[SPPM] Camera rays traced";
        tracePhotons();
        LOG(INFO) << "[SPPM] Photons traced";
        gatherResults();
    }
}

}